import pandas as pd 
import os
import numpy as np 
import warnings
warnings.filterwarnings('ignore')

# 策略逻辑 截面开仓
# 阈值以上计算排名，买入n只，开仓至金额不够

# 买入仓位表， 当日开盘价买入
# 卖出卖出表， 当日收盘价卖出
# 持仓表

# 到期平仓判断
# 开仓判断
# 止盈止损

#阈值加rank开仓方式，附带止盈止损
class Pos_make_rank_thred_step():   
    
    def __init__(self, args, open_root=None, close_root=None):
        self.args = args
        self.money_init = args.money_init
        self.money = args.money_init
        self.thred = args.open_thred
        self.vol_num = 20
        self.predict = self.get_predict()   
        self.open, self.close = self.get_price(open_root, close_root)
        self.hold_pos = self.pos_init()
        self.buy_pos = self.pos_init()
        self.sell_pos = self.pos_init()
        self.cost_df= self.pos_init()
        self.hold_ticker = {}
        self.money_df = pd.DataFrame(index=self.predict.index, columns=['money'])
        self.stop_ret = args.stop_ret
        self.stop_loss = args.stop_loss
        self.cost_ratio = args.cost_ratio
        self.hold_num = args.hold_num
        self.pos_root = self.pos_root_make()

    
    
    def pos_root_make(self):
        result_root = os.path.join(self.args.result_root, self.__class__.__name__, f'loss{self.stop_loss}_ret{self.stop_ret}_thred{self.thred}_holdnum{self.hold_num}')
        if not os.path.exists(result_root):
            os.makedirs(result_root)
        return result_root
    
    def log_write(self, content):
        with open(os.path.join(self.pos_root, 'log.txt'),"a") as file:   #只需要将之前的”w"改为“a"即可，代表追加内容
            file.write(content)

    def get_price(self, open_root, close_root):  
        # sourcery skip: avoid-builtin-shadow
        open_root = r'data/stock_data/daily/consentrate_price/open.pkl.gzip'
        close_root = r'data/stock_data/daily/consentrate_price/close.pkl.gzip'
        open = pd.read_pickle(open_root)
        close = pd.read_pickle(close_root)
        open = open.loc[self.predict.index, self.predict.columns]
        close = close.loc[self.predict.index, self.predict.columns]
        
        open = open.fillna(method='ffill')
        close = close.fillna(method='ffill')
        return open, close

    def get_predict(self):
        df = pd.read_pickle(os.path.join(self.args.result_root, 'test_pred.pkl.gzip'))
        df.index = pd.to_datetime(df.index) 
        df.dropna(how='any', axis=0, inplace=True)   
        df = df.shift(1) # 预测值滞后一天
        df = df[(df.index>=self.args.backtest_time['beg'])&(df.index<=self.args.backtest_time['end'])]
        return df
    
    def pos_init(self):
        contract = self.predict.columns
        zero_ary = np.zeros((len(self.predict),len(contract)))
        pos_df = pd.DataFrame(zero_ary, index= self.predict.index, columns=contract)
        pos_df = pos_df.astype(int)
        return pos_df
    
    def pos_make(self):         
        for idx, idv in enumerate(self.predict.index):
            # 到期平仓判断
            # self.time_stop_pos(idx, idv)

            # 开仓判断
            self.log_write(f'\n{idv}')
            print(idv)
            self.stop_ratio(idx, idv)
            self.pos_sub(idx, idv)
            self.pos_add(idx, idv)      
            self.money_df.loc[idv] = self.money
            self.hold_pos = (self.buy_pos-self.sell_pos).cumsum()
        self.pos_result_save()
    
    def pos_result_save(self):
        self.buy_pos.to_csv(f'{self.pos_root}/buy_pos.csv')
        self.sell_pos.to_csv(f'{self.pos_root}/sell_pos.csv')
        self.hold_pos.to_csv(f'{self.pos_root}/hold_pos.csv')
        self.money_df.to_csv(f'{self.pos_root}/money.csv')

    def pos_add(self, idx, idv):
        self.hold_pos = (self.buy_pos-self.sell_pos).cumsum()
        # idx_pre = idx-1
        # if idx_pre<0:
        #     asset = self.money_init
        # else:
        #     asset = (self.hold_pos.iloc[idx_pre]*self.close.iloc[idx_pre]*100).sum() + self.money_df.iloc[idx_pre]
        #     asset = asset[0]
        price_limit = 300
        pred = self.predict.loc[idv]
        open_v = self.open.loc[idv]
        open_v = open_v[(open_v >= 5) & (open_v < price_limit) ]
        pred = pred[open_v.index]
        pred = pred[pred>=0.03]
        pred = pred.sort_values(ascending=False)
        pred = pred[:self.hold_num]
        buy_codes = pred.index
        buy_dict = {code: 0 for code in buy_codes}
        price_list = self.open.loc[idv,buy_codes].values*100
        cost_lst = np.maximum(price_list*self.cost_ratio, 2)
        # open_command = any(pred[code] >= 0.8 for code in buy_codes)
        # if open_command:
        while (self.money > cost_lst+price_list).any():
            for code in buy_codes:#debug 检查此处             
                price_amount = self.open.loc[idv, code] * 100
                cost = np.maximum(price_amount*self.cost_ratio, 1)
                if self.money >= (price_amount+cost): #资金满足开仓要求
                    buy_dict[code] += 1
                    price = self.open.loc[idv, code]
                    self.money -= (price * 100) 
                    self.cost_df.loc[idv, code] += cost  
                    self.log_write(f'\n{idv}: 买入{code}, 开盘价：{(self.open.loc[idv, code]):.2f}， 花费:{price_amount:.2f}, 剩余金额：{self.money:.2f}')
                    print(f'{idv}: 买入{code}, 开盘价：{(self.open.loc[idv, code]):.2f}， 花费:{price_amount:.2f}, 剩余金额：{self.money:.2f}')

        for code, num in buy_dict.items():
            self.buy_pos.loc[idv, code] = num
            if num>0:
                if code in list(self.hold_ticker.keys()): 
                    self.hold_ticker[code][idv] = [num, self.open.loc[idv, code], self.stop_loss]
                else:
                    self.hold_ticker[code] = {}
                    self.hold_ticker[code][idv] = [num, self.open.loc[idv, code], self.stop_loss]
                


    def pos_sub(self, idx, idv):
        if idx <=0:
            return None
        self.hold_pos = (self.buy_pos-self.sell_pos).cumsum()
        pred = self.predict.loc[idv]
        pred = pred.sort_values(ascending=True)
        pred = pred[:int(self.thred*len(pred))]
        for code in list(self.hold_ticker.keys()) : 
            if code in pred.index:              
                for da in list(self.hold_ticker[code].keys()):
                    if da == idv:
                        continue
                    num = self.hold_ticker[code][da][0]
                    self.sell_pos.loc[idv, code] += num
                    price = self.open.loc[idv, code]
                    self.money += (price * 100 *num)
                    self.log_write(f'\n{idv}：{code}平仓, 剩余金额：{self.money:.2f}')
                    print(f'{idv}：{code}平仓, 剩余金额：{self.money:.2f}')
                    cost = np.maximum(price * 100 *num*self.cost_ratio, 1)
                    self.cost_df.loc[idv, code] += cost 
                    self.hold_ticker[code].pop(da) 
    
    def stop_ratio(self, idx, idv): #计算平均成本
        self.hold_pos = (self.buy_pos-self.sell_pos).cumsum()
        for code in list(self.hold_ticker.keys()):
            if not self.hold_ticker[code]:
                self.hold_ticker.pop(code)  
        for code in list(self.hold_ticker.keys()):             
            da_lst = self.hold_ticker[code].keys()
            for da in list(da_lst):
                if da == idv:
                        continue
                price_init = self.hold_ticker[code][da][1]
                price_now = self.open.loc[idv, code]
                ratio = 0
                ratio = price_now/price_init -1

                if (ratio <= self.hold_ticker[code][da][2]):          
                    num = self.hold_ticker[code][da][0]
                    if (num > 0) : #持仓大于0
                        self.sell_pos.loc[idv, code] += num
                        price = self.open.loc[idv, code]
                        self.money += (price * 100 *num)
                        self.log_write(f'\n{idv}：{code}触发次日止损平仓, 开仓日期{da}, 卖出份数：{num}， 收益率：{ratio:.2f}, 剩余金额：{self.money:.2f}')
                        print(f'{idv}：{code}触发次日止损平仓, 开仓日期{da}, 收益率：{ratio:.2f}, 剩余金额：{self.money:.2f}')
                        self.hold_ticker[code].pop(da)
                        cost = np.maximum(price * 100 *num*self.cost_ratio, 1)
                        self.cost_df.loc[idv, code] += cost
                elif (ratio >= 0) and (ratio > self.hold_ticker[code][da][2]):
                    self.hold_ticker[code][da][2] = ratio*0.8
                   
        

class Pos_make_rank_thred_month(Pos_make_rank_thred_step):

    def __init__(self, args, open_root=None, close_root=None):
        super().__init__(args, open_root, close_root)

    def get_predict(self):    
        pred_list = []
        file_list = sorted(os.listdir(self.args.result_root))
        for i in file_list:
            if i[-2:] == '01':
                df_i = pd.read_pickle(os.path.join(self.args.result_root, i, 'test_pred.pkl.gzip'))
                pred_list.append(df_i)
        df = pd.concat(pred_list, axis=0)   
        df.index = pd.to_datetime(df.index) 
        df.dropna(how='any', axis=0, inplace=True)   
        df = df.shift(1) # 预测值滞后一天
        df = df[(df.index>=self.args.backtest_time['beg'])&(df.index<=self.args.backtest_time['end'])]
        return df

    # def pos_root_make(self):
    #     result_root = os.path.join(self.args.result_root, Pos_make_rank_thred_month.__name__)
    #     return result_root


class Pos_make_series_thred_month(Pos_make_rank_thred_month):  

    def __init__(self, args, open_root=None, close_root=None):
        super().__init__(args, open_root, close_root) 
        self.predict_vol = self.predict.rolling(self.vol_num, min_periods=1).std()   
        self.vol_num = 30
        self.thred = args.open_thred

    def pos_add(self, idx, idv):
        pred = self.predict.loc[idv] / (self.predict_vol.loc[idv]+0.0000001)
        open_idv = (self.open.loc[idv]>=5)
        open_idv = open_idv[open_idv]
        pred = pred[open_idv.index]
        pred = pred.sort_values(ascending=False)

        pred = pred[:self.hold_num]
        buy_codes = pred.index
        buy_dict = {code: 0 for code in buy_codes}
        price_list = self.open.loc[idv,buy_codes].values*100
        cost_lst = np.maximum(price_list*self.cost_ratio, 2)
        while (self.money > cost_lst+price_list).any():
            for code in buy_codes:#debug 检查此处             
                price_amount = self.open.loc[idv, code] * 100
                cost = np.maximum(price_amount*self.cost_ratio, 1)
                if self.money >= (price_amount+cost): #资金满足开仓要求
                    buy_dict[code] += 1
                    price = self.open.loc[idv, code]
                    self.money -= (price * 100) 
                    self.cost_df.loc[idv, code] += cost  
                    self.log_write(f'\n{idv}: 买入{code}, 开盘价：{(self.open.loc[idv, code]):.2f}， 花费:{price_amount:.2f}, 剩余金额：{self.money:.2f}')
                    print(f'{idv}: 买入{code}, 开盘价：{(self.open.loc[idv, code]):.2f}， 花费:{price_amount:.2f}, 剩余金额：{self.money:.2f}')

        for code, num in buy_dict.items():
            self.buy_pos.loc[idv, code] = num
            if num>0:
                if code in list(self.hold_ticker.keys()): 
                    self.hold_ticker[code][idv] = [num, self.open.loc[idv, code], self.stop_loss]
                else:
                    self.hold_ticker[code] = {}
                    self.hold_ticker[code][idv] = [num, self.open.loc[idv, code], self.stop_loss]
    
    def pos_sub(self, idx, idv):
        if idx <=0:
            return None
        self.hold_pos = (self.buy_pos-self.sell_pos).cumsum()
        pred = self.predict.loc[idv] / (self.predict_vol.loc[idv]+0.0000001)
        pred = pred.sort_values(ascending=True)
        pred = pred[:int(0.05*len(pred))]
        for code in list(self.hold_ticker.keys()) : 
            if code in pred.index:              
                for da in list(self.hold_ticker[code].keys()):
                    if da == idv:
                        continue
                    num = self.hold_ticker[code][da][0]
                    self.sell_pos.loc[idv, code] += num
                    price = self.open.loc[idv, code]
                    self.money += (price * 100 *num)
                    self.log_write(f'\n{idv}：{code}平仓, 剩余金额：{self.money:.2f}')
                    print(f'{idv}：{code}平仓, 剩余金额：{self.money:.2f}')
                    cost = np.maximum(price * 100 *num*self.cost_ratio, 1)
                    self.cost_df.loc[idv, code] += cost 
                    self.hold_ticker[code].pop(da) 


class Pos_make_series_thred_step(Pos_make_series_thred_month):
    
    def __init__(self, args, open_root=None, close_root=None):
        super().__init__(args, open_root, close_root)

    def get_predict(self):
        df = pd.read_pickle(os.path.join(self.args.result_root, 'test_pred.pkl.gzip'))
        df.index = pd.to_datetime(df.index) 
        df.dropna(how='any', axis=0, inplace=True)   
        df = df.shift(1) # 预测值滞后一天
        df = df[(df.index>=self.args.backtest_time['beg'])&(df.index<=self.args.backtest_time['end'])]
        return df

    # def pos_root_make(self):
    #     result_root = os.path.join(self.args.result_root, Pos_make_series_thred_step.__name__)
    #     return result_root